import torch
import torch.nn as nn
import torch.nn.functional as F
import math


def get_torch_trans(heads=8, layers=1, channels=64):
    encoder_layer = nn.TransformerEncoderLayer(
        d_model=channels, nhead=heads, dim_feedforward=64, activation="gelu"
    )
    return nn.TransformerEncoder(encoder_layer, num_layers=layers)


def Conv1d_with_init(in_channels, out_channels, kernel_size):
    layer = nn.Conv1d(in_channels, out_channels, kernel_size)
    # Weight initialization
    nn.init.kaiming_normal_(layer.weight)
    return layer


class DiffusionEmbedding(nn.Module):
    def __init__(self, num_steps, embedding_dim=128, projection_dim=None):
        super().__init__()
        if projection_dim is None:
            projection_dim = embedding_dim
        self.register_buffer(
            "embedding",
            self._build_embedding(num_steps, embedding_dim / 2),
            persistent=False,
        )
        self.projection1 = nn.Linear(embedding_dim, projection_dim)
        self.projection2 = nn.Linear(projection_dim, projection_dim)

    def forward(self, diffusion_step):
        x = self.embedding[diffusion_step]
        x = self.projection1(x)
        x = F.silu(x)
        x = self.projection2(x)
        x = F.silu(x)
        return x

    # t_embedding(t). The embedding dimension is 128 in total for every time step t.
    def _build_embedding(self, num_steps, dim=64):
        steps = torch.arange(num_steps).unsqueeze(1)  # (T,1)
        frequencies = 10.0 ** (torch.arange(dim) / (dim - 1) * 4.0).unsqueeze(
            0
        )  # (1,dim)
        table = steps * frequencies  # (T,dim)
        table = torch.cat([torch.sin(table), torch.cos(table)], dim=1)  # (T,dim*2)
        return table

#根据当前噪声数据和条件信息，预测扩散过程中添加到原始数据的噪声（用于后续的去噪/插补）
class diff_CSDI(nn.Module):
    def __init__(self, config, inputdim=2):
        """条件扩散模型的核心网络
               Args:
                   config (dict): 配置文件，包含以下关键参数：
                       - channels: 主干网络通道数（如128）
                       - num_steps: 扩散总步数（如50）
                       - diffusion_embedding_dim: 时间步编码维度（如128）
                       - side_dim: 侧信息总维度（时间+特征+掩码）
                       - nheads: 多头注意力头数（如8）
                       - layers: 残差块层数（如6）
                       - mixed: 是否启用混合模式（bool）
                       - token_emb_dim: 混合模式下的特征嵌入维度
                   inputdim (int): 输入通道数，由CSDI_base决定：
                       - 无条件生成：1（仅噪声输入）
                       - 条件生成：2（观测值+噪声拼接）
               """
        super().__init__()
        self.config = config
        self.channels = config["channels"] # 主干网络通道数（类似UNet的基准通道数）

        # === 时间步编码器 ===
        # 将离散扩散步数t映射为连续向量（类似Transformer的位置编码）
        self.diffusion_embedding = DiffusionEmbedding(
            num_steps=config["num_steps"],# 扩散总步数T
            embedding_dim=config["diffusion_embedding_dim"], # 编码维度
        )# 输出形状：(B,) -> (B, diffusion_embedding_dim)

        # === 输入通道配置 ===
        # mixed模式允许更复杂的特征交互（如多变量联合建模）
        self.token_emb_dim = config["token_emb_dim"] if config["mixed"] else 1
        inputdim = 2 * self.token_emb_dim # 实际输入通道数（若mixed=True则扩展维度） 通常为2或2*token_emb_dim
        # === 输入投影层 ===
        # 将原始输入映射到高维空间（1x1卷积相当于全连接层）
        self.input_projection = Conv1d_with_init(inputdim, self.channels, 1) #输入通道、输出通道、1×1卷积 形状变化：(B, inputdim, K*L) -> (B, channels, K*L)
        # === 输出投影层 ===
        # 两级投影逐步降维（避免信息突变）
        self.output_projection1 = Conv1d_with_init(self.channels, self.channels, 1) # 保持通道数不变
        # 最终投影到目标维度（零初始化稳定训练初期）
        self.output_projection2 = Conv1d_with_init(self.channels, self.token_emb_dim, 1)
        nn.init.zeros_(self.output_projection2.weight) # 初始预测接近零噪声

        # === 残差块堆叠 ===
        self.residual_layers = nn.ModuleList(
            [
                ResidualBlock(
                    side_dim=config["side_dim"], # 侧信息总维度
                    channels=self.channels,# 特征通道数
                    diffusion_embedding_dim=config["diffusion_embedding_dim"], # 时间编码维度
                    nheads=config["nheads"], # 注意力头数
                )
                for _ in range(config["layers"]) # 堆叠层数（如6层）
            ]
        )

    def forward(self, x, cond_info, diffusion_step):
        """前向传播
               Args:
                   x: 模型输入
                       - 无条件：(B, 1, K, L) 仅噪声
                       - 条件：(B, 2, K, L) [观测值, 噪声]
                   cond_info: 侧信息 (B, side_dim, K, L)
                   diffusion_step: 当前扩散步数 (B,)
               Returns:
                   预测的噪声 (B, K, L)
               """
        B, inputdim, K, L = x.shape# Batch, 输入通道, 特征数, 序列长度
        # === 输入预处理 ===
        # 展平空间维度：将特征(K)和时间(L)合并，使1D卷积能同时处理时空信息
        x = x.reshape(B, inputdim, K * L)
        # 通道投影：将输入映射到高维空间
        x = self.input_projection(x)
        x = F.relu(x) # 激活函数引入非线性
        x = x.reshape(B, self.channels, K, L)  # 恢复空间结构：分离特征和时间维度
        # === 时间步编码 ===
        diffusion_emb = self.diffusion_embedding(diffusion_step)# (B,) -> (B, diffusion_embedding_dim）
        # === 残差块处理 ===
        skip = []# 存储各层跳跃连接
        for layer in self.residual_layers:
            # 每个残差块处理流程：
            # 1. 时空注意力（捕捉远程依赖）
            # 2. 条件层归一化（融合扩散步编码和侧信息）
            # 3. 门控卷积（局部特征提取）
            x, skip_connection = layer(x, cond_info, diffusion_emb)
            skip.append(skip_connection)# 收集跳跃输出
        # 跳跃连接聚合：类似UNet的跳连机制
        # 各层跳跃连接求和后归一化（稳定数值范围）
        x = torch.sum(torch.stack(skip), dim=0) / math.sqrt(len(self.residual_layers))
        # (B, channels, K, L)
        # === 输出处理 ===
        # 再次展平空间维度
        x = x.reshape(B, self.channels, K * L)
        # 第一级投影：保持通道数不变
        x = self.output_projection1(x)# (B, channels, K*L)
        x = F.relu(x)
        # 第二级投影：降维到目标输出形状
        x = self.output_projection2(x) # (B, token_emb_dim, K*L)
        # === 输出形状调整 ===
        if self.config["mixed"]:
            # 混合模式：允许特征维度交互
            x = x.permute(0, 2, 1)# (B, K*L, token_emb_dim)
            x = x.reshape(B, K, L * self.token_emb_dim) # 特征混合
        else:
            x = x.reshape(B, K, L)# 常规模式：独立预测每个特征 (B, K, L)
        return x # 输出预测的噪声


class ResidualBlock(nn.Module):
    def __init__(self, side_dim, channels, diffusion_embedding_dim, nheads):
        super().__init__()
        self.diffusion_projection = nn.Linear(diffusion_embedding_dim, channels)
        self.cond_projection = Conv1d_with_init(side_dim, 2 * channels, 1)
        self.mid_projection = Conv1d_with_init(channels, 2 * channels, 1)
        self.output_projection = Conv1d_with_init(channels, 2 * channels, 1)

        # Temporal Transformer layer
        self.time_layer = get_torch_trans(heads=nheads, layers=1, channels=channels)
        # Feature Transformer layer
        self.feature_layer = get_torch_trans(heads=nheads, layers=1, channels=channels)

    def forward_time(self, y, base_shape):
        B, channel, K, L = base_shape
        if L == 1:
            return y
        y = y.reshape(B, channel, K, L).permute(0, 2, 1, 3).reshape(B * K, channel, L)
        # (B*K, C, L) -> (L, B*K, C) -> (B*K, C, L)
        # input shape for transformerencoder: [seq, batch, emb]
        y = self.time_layer(y.permute(2, 0, 1)).permute(1, 2, 0)
        y = y.reshape(B, K, channel, L).permute(0, 2, 1, 3).reshape(B, channel, K * L)
        return y

    def forward_feature(self, y, base_shape):
        B, channel, K, L = base_shape
        if K == 1:
            return y
        y = y.reshape(B, channel, K, L).permute(0, 3, 1, 2).reshape(B * L, channel, K)
        # (B*L, C, K) -> (K, B*L, C) -> (B*L, C, K)
        y = self.feature_layer(y.permute(2, 0, 1)).permute(1, 2, 0)
        y = y.reshape(B, L, channel, K).permute(0, 2, 3, 1).reshape(B, channel, K * L)
        return y

    def forward(self, x, cond_info, diffusion_emb):
        B, channel, K, L = x.shape
        base_shape = x.shape
        x = x.reshape(B, channel, K * L)

        # diffusion_emb is
        diffusion_emb = self.diffusion_projection(diffusion_emb).unsqueeze(
            -1
        )  # (B,channel,1)
        y = x + diffusion_emb

        y = self.forward_time(y, base_shape)
        y = self.forward_feature(y, base_shape)  # (B,channel,K*L)
        y = self.mid_projection(y)  # (B,2*channel,K*L)

        _, cond_dim, _, _ = cond_info.shape
        cond_info = cond_info.reshape(B, cond_dim, K * L)
        cond_info = self.cond_projection(cond_info)  # (B,2*channel,K*L)
        y = y + cond_info

        gate, filter = torch.chunk(y, 2, dim=1)
        y = torch.sigmoid(gate) * torch.tanh(filter)  # (B,channel,K*L)
        y = self.output_projection(y)

        residual, skip = torch.chunk(y, 2, dim=1)
        x = x.reshape(base_shape)
        residual = residual.reshape(base_shape)
        skip = skip.reshape(base_shape)
        return (x + residual) / math.sqrt(2.0), skip
